package ws.very.util.lang.macroPack

import scala.language.experimental.macros

trait BlockChildrenValCollecter {
  //def collect[T](collect: PartialFunction[Any, T])(block: Unit): Seq[T] = macro collectImpl[T]

  def collectImpl[T: c.WeakTypeTag](c: scala.reflect.macros.Context)(collect: c.Expr[PartialFunction[Any, T]])(block: c.Expr[Unit]): c.Expr[Seq[T]] = {
    import c.universe._

    def wrap(tree: Tree) = Apply(Select(Ident(newTermName("BlockChildrenValCollect")), newTermName("apply")), tree :: Nil)

    def lineChange(line: Tree): List[Tree] = line match {
      case line: Apply => wrap(line) :: Nil
      case line: Select => wrap(line) :: Nil
      case line: Literal => wrap(line) :: Nil
      case line @ ValDef(modifiers, termName, typeTree, value) => ValDef(modifiers, termName, typeTree, lineChange(value).head) :: Nil
      //case line @ ModuleDef(modifiers,termName,tpls)=>line::wrap(Ident(termName))::Nil
      case Block(stats, expr) => Block(stats.flatMap { line => lineChange(line) }, lineChange(expr).head) :: Nil
      case line => line :: Nil

    }
    val newBlock = c.Expr[Unit](lineChange(block.tree).head)
    c.Expr[Seq[T]](reify {
      object BlockChildrenValCollect {
        var result = scala.collection.mutable.Buffer[T]()
        def apply[T](t: T) = {
          Option(t).collect(collect.splice).foreach { t => result += t }
          t
        }
      }

      newBlock.splice

      BlockChildrenValCollect.result
    }.tree)
  }

}

object BlockChildrenValCollecters extends BlockChildrenValCollecter
